"""
20_capstone_project.py - Full-Stack AI Agent System
=====================================================
Chapter 20: Putting everything together.

ARCHITECTURE:
    User Query → Supervisor → Specialist Agent → Synthesizer → Response

Features:
  - Supervisor routing to 4 specialist agents
  - RAG, Code, Web, and General agents
  - Memory with checkpointer
  - Execution tracing
  - Error handling

Run: python src/20_capstone_project.py
"""

from typing import TypedDict, Annotated, Literal
from operator import add
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver


class AgentState(TypedDict):
    """Complete agent system state."""
    query: str
    messages: Annotated[list, add]
    active_agent: str
    context: str
    tool_outputs: Annotated[list, add]
    final_response: str
    iteration: int


def supervisor(state: AgentState) -> dict:
    """Route to the best specialist agent."""
    query = state["query"].lower()
    
    if any(w in query for w in ["search", "find", "what is", "explain"]):
        agent = "rag_agent"
    elif any(w in query for w in ["code", "python", "function", "script"]):
        agent = "code_agent"
    elif any(w in query for w in ["news", "latest", "today", "current"]):
        agent = "web_agent"
    else:
        agent = "general_agent"
    
    return {
        "active_agent": agent,
        "messages": [f"[Supervisor] Routing to {agent}"],
        "iteration": state.get("iteration", 0) + 1
    }


def rag_agent(state: AgentState) -> dict:
    """Retrieval-Augmented Generation agent."""
    return {
        "context": f"Retrieved 3 relevant documents for: {state['query']}",
        "messages": ["[RAG Agent] Vector search complete. Context retrieved."],
        "tool_outputs": [{"agent": "rag", "status": "success", "docs": 3}]
    }


def code_agent(state: AgentState) -> dict:
    """Code generation agent."""
    return {
        "context": f"Generated Python code for: {state['query']}",
        "messages": ["[Code Agent] Code generated and validated."],
        "tool_outputs": [{"agent": "code", "status": "success"}]
    }


def web_agent(state: AgentState) -> dict:
    """Web search agent."""
    return {
        "context": f"Found 5 web results for: {state['query']}",
        "messages": ["[Web Agent] Web search complete."],
        "tool_outputs": [{"agent": "web", "status": "success", "results": 5}]
    }


def general_agent(state: AgentState) -> dict:
    """General conversation agent."""
    return {
        "context": f"General knowledge response for: {state['query']}",
        "messages": ["[General Agent] Response prepared."],
        "tool_outputs": [{"agent": "general", "status": "success"}]
    }


def synthesizer(state: AgentState) -> dict:
    """Combine all agent outputs into final response."""
    agent = state["active_agent"]
    context = state.get("context", "No context available")
    
    response = (
        f"Based on {agent.replace('_', ' ').title()}'s analysis:\n"
        f"{context}\n"
        f"[Processed in {state.get('iteration', 1)} iteration(s)]"
    )
    
    return {
        "final_response": response,
        "messages": ["[Synthesizer] Final response assembled."]
    }


def route_to_agent(state: AgentState) -> str:
    return state["active_agent"]


# ============================================================
# BUILD THE COMPLETE SYSTEM
# ============================================================

graph = StateGraph(AgentState)

graph.add_node("supervisor", supervisor)
graph.add_node("rag_agent", rag_agent)
graph.add_node("code_agent", code_agent)
graph.add_node("web_agent", web_agent)
graph.add_node("general_agent", general_agent)
graph.add_node("synthesizer", synthesizer)

graph.add_edge(START, "supervisor")
graph.add_conditional_edges("supervisor", route_to_agent)

for agent in ["rag_agent", "code_agent", "web_agent", "general_agent"]:
    graph.add_edge(agent, "synthesizer")

graph.add_edge("synthesizer", END)

checkpointer = MemorySaver()
app = graph.compile(checkpointer=checkpointer)


if __name__ == "__main__":
    print("=" * 60)
    print("20: Capstone - Full-Stack AI Agent System")
    print("=" * 60)
    
    test_queries = [
        "What is retrieval augmented generation?",
        "Write a Python function for binary search",
        "What are the latest AI news today?",
        "Tell me a joke about programming",
    ]
    
    for i, query in enumerate(test_queries):
        config = {"configurable": {"thread_id": f"capstone_{i}"}}
        
        print(f"\n{'─' * 50}")
        print(f"Query: {query}")
        print(f"{'─' * 50}")
        
        result = app.invoke({"query": query}, config=config)
        
        print("\n  Execution Trace:")
        for msg in result["messages"]:
            print(f"    {msg}")
        
        print(f"\n  📝 Response:\n    {result['final_response']}")
    
    print(f"\n{'=' * 60}")
    print("✅ Capstone project complete!")
    print("   All 20 chapters implemented!")
    print("=" * 60)
